package org.paraj.mapreduce.executors;

import com.google.inject.Inject;
import org.paraj.mapreduce.ComplexTask;
import org.paraj.mapreduce.ResultAggregator;
import org.paraj.mapreduce.SimpleTask;
import org.paraj.mapreduce.Task;

import java.io.Serializable;
import java.rmi.RemoteException;
import java.util.Collection;
import java.util.UUID;

public class MultithreadedTaskExecutor implements MapReduceExecutor{

    private Distributor distributor;
    ResultService resultService;

    public  MultithreadedTaskExecutor(final ResultService resultService, final Distributor distributor) throws RemoteException {
        this.resultService = resultService;
        this.distributor = distributor;
    }

    public <R> R run(final Task<R> task) throws InterruptedException{
        final UUID uid = scheduleTask(null, task);
        return (R)resultService.waitForResult(uid);
    }


    public static  class RunTask implements Runnable, Serializable{
        private final Task task;
        private UUID parentId;
        private UUID uid;
        @Inject transient private ResultService resultService;
        @Inject transient private Distributor distributor;

        public RunTask(UUID parentId, UUID uid, Task task){
            this.uid = uid;
            this.parentId = parentId;
            this.task = task;
        }


        public void run(){
            if (task instanceof SimpleTask) {
                resultService.finishedTask(uid, ((SimpleTask)task).execute());
            }

            if (task instanceof ComplexTask) {
                executeComplexTask((ComplexTask) task);
            }
        }

        private void executeComplexTask(ComplexTask task) {
            final Collection<Task> subtasks = task.generateSubtasks();

            resultService.initializeComplexTask( uid, subtasks.size());
            if (task instanceof StatefulTask){
                scheduleSubtasks(uid, subtasks);
                resultService.finishedSubtask(parentId, aggregate(uid, task.getResultAggregator()));
            }else{
                resultService.registerAggregator(uid, task.getResultAggregator());
                scheduleSubtasks(uid, subtasks);
            }
        }

        private Object aggregate(UUID uid, ResultAggregator resultAggregator) {
            try {
                return resultAggregator.aggregate(resultService.waitForSubResults(uid));
            } catch (InterruptedException e) {
                throw new RuntimeException("Interrupted while waiting for subtasks results", e);
            }
        }

        private void scheduleSubtasks(UUID uid, Collection<Task> subtasks) {
            for(Task subtask: subtasks){
                scheduleTask(uid, subtask);
            }
        }

        public <R> UUID scheduleTask(final UUID parentId, Task<R> task) {
            final UUID uid = UUID.randomUUID();
            resultService.startingTask(parentId, uid);

            distributor.schedule(new MultithreadedTaskExecutor.RunTask(parentId, uid, task));
            return uid;
        }
        
        public void setResultService(ResultService resultService) {
            this.resultService = resultService;
        }

        public void setDistributor(Distributor distributor) {
            this.distributor = distributor;
        }

        @Override
        public String toString() {
            return "RunTask{" +
                    "uid=" + uid +
                    ",task=" + task +
                    ", parentId=" + parentId +
                    '}';
        }
    }


    public <R> UUID scheduleTask(final UUID parentId, Task<R> task) {
        final UUID uid = UUID.randomUUID();
        resultService.startingTask(parentId, uid);

        distributor.schedule(new MultithreadedTaskExecutor.RunTask(parentId, uid, task));
        return uid;
    }

}
